from util_flux import process_img_1024
from preprocess.humanparsing.run_parsing import Parsing
from PIL import Image
import numpy as np
import cv2
import os

parse_ckp = '/mnt/nas/shengjie/huggingface_model_local/ootd/checkpoints/humanparsing'
target_shape = (384, 512)

label_map = {
    "background": 0,
    "hat": 1,
    "hair": 2,
    "sunglasses": 3,
    "upper_clothes": 4,
    "skirt": 5,
    "pants": 6,
    "dress": 7,
    "belt": 8,
    "left_shoe": 9,
    "right_shoe": 10,
    "head": 11,
    "left_leg": 12,
    "right_leg": 13,
    "left_arm": 14,
    "right_arm": 15,
    "bag": 16,
    "scarf": 17,
    "neck": 18,
}

upper_zone = [
    "upper_clothes",
    "left_arm", "right_arm"
]
lower_zone = [
    "skirt", "pants","belt",
    "left_leg", "right_leg",
]
full_zone = ["dress"] + upper_zone + lower_zone

required_zone = [
    'neck',
    'hat','hair','sunglasses',
    'left_shoe','right_shoe',
    'head','bag',
    'scarf',
]

def get_parse_predictor():
    # 检查模型路径是否存在
    if not os.path.exists(parse_ckp):
        raise FileNotFoundError(f"Parsing checkpoint path does not exist: {parse_ckp}")
    parsing_model = Parsing(0, prefix=parse_ckp)
    return parsing_model

def get_parse_res(parse_predictor=None, img_path='', img_pil=None, target_shape=(384, 512)):
    if img_pil is None:
        if not img_path or not os.path.isfile(img_path):
            raise ValueError("Either img_pil must be provided or img_path must be a valid file path.")
        # img_pil = Image.open(img_path).convert("RGB")
        img_pil = process_img_1024(img_path,target_shape=target_shape).convert("RGB")
    else:
        if not isinstance(img_pil, Image.Image):
            raise TypeError("img_pil must be a PIL.Image.Image object.")
        img_pil = process_img_1024('',img_pil,target_shape=target_shape)

    model_img_pil = img_pil.resize(target_shape, Image.BILINEAR)
    if parse_predictor is None:
        raise ValueError("parse_predictor must be provided.")
    model_parse, _ = parse_predictor(model_img_pil)
    return model_parse

def get_parse_zone_by_clotype(clo_type='', parse_pil=None, target_shape=(384, 512)):
    """
    根据clo_type选择不同的zone，并返回对应的mask（PIL Image）。
    clo_type: str, 可选 'upper', 'lower', 'full'
    parse_pil: PIL Image, 解析图像（单通道，每个像素为label id）
    返回: PIL Image, mask（0/255）
    """
    valid_types = {'upper', 'lower', 'full'}
    if clo_type not in valid_types:
        raise ValueError(f"clo_type must be one of {valid_types}, but got '{clo_type}'")
    if parse_pil is None:
        raise ValueError("parse_pil must be provided.")

    # Resize and convert to numpy array
    im_parse = parse_pil.resize(target_shape, Image.NEAREST)
    parse_array = np.array(im_parse)
    mask = np.zeros_like(parse_array, dtype=np.uint8)

    # 选择zone
    if clo_type == 'upper':
        zone_labels = upper_zone
    elif clo_type == 'lower':
        zone_labels = lower_zone
    elif clo_type == 'full':
        zone_labels = full_zone
    else:
        # 默认返回全0 mask
        return Image.new('L', parse_pil.size, 0)

    global required_zone

    # Build masks by aggregating relevant labels
    # 先加zone_labels
    for label_name in zone_labels:
        label_value = label_map.get(label_name, None)
        if label_value is not None:
            mask = np.logical_or(mask, (parse_array == label_value)).astype(np.uint8)
    # 再减去required_zone（如果有定义）
    assert required_zone is not None, required_zone
    # print('required_zone', required_zone)
    for label_name in required_zone:
        label_value = label_map.get(label_name, None)
        if label_value is not None:
            mask = np.logical_and(mask, (parse_array != label_value)).astype(np.uint8)

    # Apply morphological operations
    kernel = np.ones((3, 3), np.uint8)

    def refine_mask(mask):
        """Apply morphological refinement to mask"""
        mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)
        mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)
        return mask

    # Refine mask
    mask_refined = refine_mask(mask)

    # Convert to PIL Image
    parse_zone = Image.fromarray((mask_refined * 255).astype(np.uint8), mode='L')

    return parse_zone

if __name__ == '__main__':
    parse_predictor = get_parse_predictor()

    img_path = './tmp_tryon.jpg'
    if not os.path.isfile(img_path):
        raise FileNotFoundError(f"Image file not found: {img_path}")

    parse_pil = get_parse_res(parse_predictor, img_path=img_path, target_shape=target_shape)

    # 测试三种clo_type
    for clo_type in ['upper', 'lower', 'full']:
        mask = get_parse_zone_by_clotype(clo_type=clo_type, parse_pil=parse_pil, target_shape=target_shape)
        out_path = f'tmp_tryon2_{clo_type}.jpg'
        mask.save(out_path)
        print(f"Saved mask for {clo_type} to {out_path}")